package Structures;

/**
 * 前缀树，适用于全部由小写字母构成的字符串
 */
public class Trie {
    class TrieNode {
        int endCount = 0; // 有多少单词以本结点结尾
        int prefixCount = 0; // 有多少单词以包含此结点在内的字符串为前缀
        TrieNode[] children;

        public TrieNode(int nChildren) {
            children = new TrieNode[nChildren];
        }
    }

    TrieNode root = new TrieNode(26);

    public Trie() {
    }

    public void insert(String word) {
        char[] charArray = word.toCharArray();
        TrieNode ptr = root;
        for (int i = 0; i < charArray.length; i++) {
            char ch = charArray[i];
            if (ptr.children[ch - 'a'] == null) ptr.children[ch - 'a'] = new TrieNode(26);
            ptr = ptr.children[ch - 'a'];
            ptr.prefixCount++;
        }

        ptr.endCount++;
    }

    /**
     * 检查前缀树中是否有字符串word
     */
    public boolean search(String word) {
        if (word.isEmpty()) return false;

        char[] charArray = word.toCharArray();
        TrieNode ptr = root;
        for (int i = 0; i < charArray.length; i++) {
            char ch = charArray[i];
            if (ptr.children[ch - 'a'] == null) return false;
            ptr = ptr.children[ch - 'a'];
        }
        return ptr.endCount > 0;
    }

    /**
     * 检查前缀树中是否有字符串以prefix开头
     */
    public boolean startsWith(String prefix) {
        if (prefix.isEmpty()) return false;

        char[] charArray = prefix.toCharArray();
        TrieNode ptr = root;
        for (int i = 0; i < charArray.length; i++) {
            char ch = charArray[i];
            if (ptr.children[ch - 'a'] == null) return false;
            ptr = ptr.children[ch - 'a'];
        }
        return true;
    }
}
